import paddle
import paddle.nn as nn
import paddle.nn.functional as F

class FeedForward(nn.Layer):
    def __init__(self, hidden_dim):
        super(FeedForward, self).__init__()
        self.fc_one = nn.Linear(hidden_dim, hidden_dim // 2, bias_attr=False)
        self.fc_two = nn.Linear(hidden_dim // 2, hidden_dim, bias_attr=False)
        self.gelu = nn.GELU()

    def forward(self, feed_x):
        feed_x = self.fc_one(feed_x)
        feed_x = self.gelu(feed_x)
        feed_x = self.fc_two(feed_x)
        return feed_x

class Block(nn.Layer):
    def __init__(self, hidden_dim):
        super(Block, self).__init__()
        self.layer_nor_x = nn.LayerNorm(hidden_dim)
        self.layer_nor_p = nn.LayerNorm(hidden_dim)
        self.layer_a = nn.Linear(2 * hidden_dim, hidden_dim, bias_attr=True)
        self.layer_b = nn.Linear(2 * hidden_dim, hidden_dim, bias_attr=True)
        self.layer_c = nn.Linear(2 * hidden_dim, hidden_dim, bias_attr=True)
        self.layer_abc = nn.Linear(2 * hidden_dim, hidden_dim, bias_attr=True)

        self.layer_nor_a = nn.LayerNorm(hidden_dim)
        self.layer_nor_b = nn.LayerNorm(hidden_dim)
        self.layer_nor_c = nn.LayerNorm(hidden_dim)
        self.mlp = FeedForward(hidden_dim)

    def forward(self, p, x):
        x = self.layer_nor_x(x)
        p = self.layer_nor_p(p)
        a = self.layer_a(paddle.concat([x, p], -1))
        b = self.layer_b(paddle.concat([x, p], -1))
        c = self.layer_c(paddle.concat([x, p], -1))
        a = paddle.cumsum(a, 1)
        b = paddle.cumsum(b, 1)
        c = paddle.cumsum(c, 1)
        a = self.layer_nor_a(a)
        b = self.layer_nor_b(b)
        c = self.layer_nor_c(c)

        abc = paddle.concat([a.unsqueeze(-1), b.unsqueeze(-1), c.unsqueeze(-1)], -1)
        x = paddle.concat([F.max_pool1d(abc.reshape([abc.shape[0],-1,3]), 3).reshape(x.shape), x], -1)
        x = self.layer_nor_a(self.layer_abc(x))

        x = self.layer_nor_b(self.mlp(x) + x)
        return x

class CvFo(nn.Layer):
    def __init__(self, voc_size, hidden_dim, row_layers):
        super(CvFo, self).__init__()
        self.em = nn.Embedding(voc_size, hidden_dim)
        self.pos = nn.Linear(hidden_dim, hidden_dim, bias_attr=False)
        self.block = nn.LayerList([Block(hidden_dim) for _ in range(row_layers - row_layers // 3)])
        self.q_star = nn.LayerList([nn.Linear(2 * hidden_dim, hidden_dim, bias_attr=False) for _ in range(row_layers // 3)])

        self.out_layer = nn.Linear(hidden_dim, voc_size, bias_attr=False)

    def forward(self, sx):
        sx = self.em(sx)
        q_star = sx
        psx = self.pos(sx)
        for block in self.block:
            sx = block(psx, sx)
        for star in self.q_star:
            sx = star(paddle.concat([sx, q_star], -1))

        out = self.out_layer(sx)
        return out

if __name__ == '__main__':
    net = CvFo(512, 412, 8)
    net(paddle.randint(0, 123, [2, 15]))
